# from importlib import import_module
from flask import Flask, jsonify, request

# import numpy as np
# import torch
from infer import classify_text

app = Flask(__name__)

# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# torch.cuda.set_device(1)
# # 加载模型
# dataset = 'TikTok'  # 数据集
# model_name = 'ERNIE'  # bert
# x = import_module(f'models.{model_name}')
# config = x.Config(dataset)
# np.random.seed(3407)
# torch.manual_seed(3407)
# torch.cuda.manual_seed_all(3407)
# torch.backends.cudnn.deterministic = True  # 保证每次结果一样
# model = x.Model(config).to(device)
# model.load_state_dict(torch.load(config.save_path))
# model.eval()


@app.route('/')
def hello():
    return 'Hello World!'


@app.route('/predict', methods=['POST'])
def predict():
    # 从请求中获取参数
    if 'video' not in request.json:
        return jsonify({"error": "Missing 'video' key in JSON request"}), 400

    videos = request.json['video']
    results = []

    for video in videos:
        video_id = video.get('video_id')
        title = video.get('title')

        if not video_id or not title:
            results.append({
                'video_id': video_id,
                'error': 'Missing video_id or title'
            })
            continue
        prediction = classify_text(title)
        # prediction_class = "fake prediction for demonstration"
        results.append({
            'video_id': video_id,
            'prediction': prediction
        })

    return jsonify(results)


if __name__ == '__main__':
    app.run(host='0.0.0.0', debug=True, port=5001)
